#-*- coding=utf-8 -*-
import partie1 as p1
import numpy as np
import matplotlib.pyplot as plt

def plot_sol_2(t0,y0,f):
    "Affiche la solution approximée sur l' intervalle [0 10]"
    nmax = 200
    x = np.arange(0.,10.+10./nmax,10./nmax)
    res = p1.meth_n_step(y0, t0, nmax, 10./nmax, f, p1.step_kutta)[1]
    return plt.plot(x, res, linewidth=1.0)

# ************* MALTHUS MODEL *******************
# First modelisation with positive coefficient
plt.clf()
alpha = 2.8  # growing rate
beta = 1.9  # death rate
t0 = np.array([[0.]])
y0 = np.array([[6e9]]) # initial value
f = lambda X, T: np.array([[(alpha-beta)*X[0][0]]])
g1=plot_sol_2(t0,y0,f)
plt.legend((g1,), ("a-b=0.9",))
plt.show()

# Second modelisation with positive coefficient
plt.clf()
alpha = 2.8  # growing rate
beta = 3.2  # death rate
t0 = np.array([[0.]])
y0 = np.array([[6e9]]) # initial value
f = lambda X, T: np.array([[(alpha-beta)*X[0][0]]])
g2=plot_sol_2(t0,y0,f)
plt.legend((g2,), ("a-b=-0.4",))
plt.show()


# *********** VERHULST MODEL ******************
# First modelisation for the human condition :
plt.clf()
alpha = 2.8  # growing rate
beta = 20e9  # carrying capacity
t0 = np.array([[0.]])
y0 = np.array([[6e9]]) # initial value
f = lambda X, T: np.array([[alpha*X[0][0]*(1 - X[0][0]/beta)]])
g1=plot_sol_2(t0,y0,f)

# Second modelisation for the human condition :
alpha = 2.8  # growing rate
beta = 10e9  # carrying capacity
t0 = np.array([[0.]])
y0 = np.array([[6e9]]) # initial value
f = lambda X, T: np.array([[alpha*X[0][0]*(1 - X[0][0]/beta)]])
g2=plot_sol_2(t0,y0,f)

plt.legend((g1,g2),("kappa = 20e9","kappa = 10e9"))

plt.show()



#************* Lotka - Volterra ******************
def plot_sol_2d(t0,y0,f, colors = np.array([[0.5,0.,0.],[0.,0.5,0.]])):
    "plot in 2 dimension"
    nmax = 200
    interval_lentgh = 40.
    x = []
    y = []
    res = p1.meth_n_step(y0, t0, nmax, interval_lentgh/nmax, f, p1.step_kutta)[0]
    for l in range(len(res)):
        x.append([res[l][0][0]])
        y.append([res[l][1][0]])
    xaxis = np.arange(0., interval_lentgh + interval_lentgh/nmax, interval_lentgh/nmax)
    if(colors[1][0] == 1. and colors[1][1] == 1.):
        return plt.plot(xaxis, x, color=(colors[0][0],colors[0][1],colors[0][2]))
    return (plt.plot(xaxis, x, color=(colors[0][0],colors[0][1],colors[0][2])), plt.plot(xaxis, y, color=(colors[1][0],colors[1][1],colors[1][2])))


def freq_measure(t0,y0,f):
    "measure frenquency of the f-curve"
    nmax = 200
    interval_lentgh = 100.
    x = []
    y = []
    res = p1.meth_n_step(y0, t0, nmax, interval_lentgh/nmax, f, p1.step_kutta)[0]  # n values between 0. and 100.
    # counts the number of y0 between 0. and 100. :
    j = -1.
    mem = y0[1][0]
    for i in range(len(res)):
        # just working on the number of predators (same period as the number of predators)
        if(mem <= y0[1][0] and res[i][1] >= y0[1][0]): 
            j += 1.
        mem = res[i][1]
    if (j == 0.):
        return 0.
    return 100./j

# Initialisation taken from :
alpha = 0.2   # natural growth of preys in abscence of predator
beta = 0.4    # natural death rate of predators in abscence of food
gamma = 0.7   # death rate per encounter of preys due to predation
delta = 0.5   # natural rate of efficient hunter predator

# Double of predators :
t0 = np.array([[0.]])
y0 = np.array([[0.5],
               [1.]])
f = lambda X,T: np.array([[X[0][0] * (alpha - beta * X[1][0])],
                           [X[1][0] * (delta * X[0][0] - gamma)]])
(g1,g2) = plot_sol_2d(t0,y0,f)
plt.legend((g1,g2),("Proies","Predateurs"))
print "frequency measure of the last curve : ", freq_measure(t0,y0,f)
plt.show()
plt.clf()
print "ploting the evolution of the prey number curve on a little square..."
for i in range(10):
    for j in range(10):
        y0 = np.array([[0.5 + 0.05*i],  # initial number of preys
                       [1. - 0.05*j]]) # initial number of predators
        plot_sol_2d(t0,y0,f, np.array([[float(i)/10,float(j)/10,0.],[1.,1.,1.]]))
print "... done !"
plt.show()

# Looking for a constant solution
t0 = np.array([[0.]])
y0 = np.array([[gamma/delta],  # initial number of preys
               [alpha/beta]]) # initial number of predators
f = lambda X,T: np.array([[X[0][0] * (alpha - beta * X[1][0])],
                           [X[1][0] * (delta * X[0][0] - gamma)]])
plt.axis([0.,40.,0.,2.])
(g1,g2) = plot_sol_2d(t0,y0,f)
plt.legend((g1,g2),("Proies","Predateurs"))
plt.show()


def plot_sol_2d_xy(t0,y0,f, colors=np.array([[1.],[1.],[1.]])):
    x = []
    y = []
    nmax = 100
    interval_length = 40.
    res = p1.meth_n_step(y0, t0,nmax, interval_length/nmax, f, p1.step_kutta)[0]
    for l in range(len(res)):
        x = x + [res[l][0][0]]
        y = y + [res[l][1][0]]
    if(colors[0][0] == 1. and colors[1][0] == 1.):
        return plt.plot(x,y,'b')
    else:
        return plt.plot(x,y,color=(colors[0][0],colors[1][0],0.))

# periodic initialisation (double of predators)
t0 = np.array([[0.]])
f = lambda X,T: np.array([[X[0][0] * (alpha - beta * X[1][0])],
                           [X[1][0] * (delta * X[0][0] - gamma)]])
N = 10

for i in range(N+1):
    for j in range(N+1):
        y0 = np.array([[0. + 0.1 * i],  # initial number of preys
                       [0. + 0.1 * j]])  # initial number of predators
        plot_sol_2d_xy(t0,y0,f, np.array([[float(i)/N],[float(j)/N],[0.]]))
plt.axis([0.,10.,0.,5.])
plt.show()

# print "simple representation in order to show a linear evolution to the equilibrum point"
for i in range(N+1):
    y0 = np.array([[0. + 0.1 * i],  # initial number of preys
                   [0. + 0.1 * i]])  # initial number of predators
    plot_sol_2d_xy(t0,y0,f, np.array([[float(i)/N],[0.],[0.]]))
plt.axis([0.,10.,0.,5.])
plt.show()
